import torch
import time
from torch.optim.lr_scheduler import StepLR  # 可选：学习率调度器


def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs, device, save_path):
    best_accuracy = 0

    # 将模型移动到指定的设备（GPU 或 CPU）
    model.to(device)

    # 可选：定义学习率调度器
    scheduler = StepLR(optimizer, step_size=10, gamma=0.1)  # 每10个epoch衰减一次学习率

    for epoch in range(num_epochs):
        start_time = time.time()  # 记录epoch开始时间

        # 训练阶段
        model.train()
        total_loss = 0
        for inputs, targets in train_loader:
            inputs, targets = inputs.to(device), targets.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # 验证阶段
        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for inputs, targets in test_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs, 1)
                total += targets.size(0)
                correct += (predicted == targets).sum().item()

        accuracy = correct / total
        print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {total_loss / len(train_loader):.4f}, Accuracy: {accuracy:.4f}")

        # 计算并打印耗时
        epoch_time = time.time() - start_time  # 计算epoch耗时
        print(f"Epoch {epoch + 1} took {epoch_time:.2f} seconds.")

        # 保存最佳模型
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.state_dict(), save_path)
            print(f"Best model saved with accuracy: {accuracy:.4f}")

        # 可选：学习率调度
        # scheduler.step()

